package asdf.springmvc;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;

/**
 * Created by chengwenlee on 12/11/15.
 */
public class ResponseWrapper extends HttpServletResponseWrapper {

    private YungeServletOutputStream outputStream;

    public ResponseWrapper(HttpServletResponse response) {
        super(response);
        outputStream = new YungeServletOutputStream(response);
    }

    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        return this.outputStream;
    }

    private static class YungeServletOutputStream extends ServletOutputStream {

        private static final Logger LGR = LoggerFactory.getLogger(YungeServletOutputStream.class);

        private ByteArrayOutputStream intermediary;
        private ServletOutputStream realOne;
        private volatile boolean flushed;

        YungeServletOutputStream(HttpServletResponse response) {
            try {
                this.realOne = response.getOutputStream();
                intermediary = new ByteArrayOutputStream();
            } catch (IOException e) {
                LGR.error("failed to get ServletOutputStream from response.", e);
            }
        }

        @Override
        public void write(int b) throws IOException {
            intermediary.write(b);
        }

        @Override
        public void flush() throws IOException {
            if (flushed) {
                return;
            }
            flushed = true;
            String out = intermediary.toString();
            LGR.debug("content to encrypt:" + out);
            String responseText = CryptUtils.encryptText(out);
            realOne.print(responseText);
            realOne.flush();
        }

        @Override
        public void close() throws IOException {
            this.flush();
            this.realOne.close();
        }

        @Override
        public boolean isReady() {
            return true;
        }

        @Override
        public void setWriteListener(WriteListener writeListener) {

        }
    }

}
